// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/file/file_input_stream.h>
#include <pollux/common/file/file_systems.h>
#include <pollux/exec/split.h>
#include <pollux/exec/trace.h>
#include <pollux/serializers/presto_serializer.h>

namespace kumo::pollux::exec::trace {
    /// Used to read an operator trace input.
    class OperatorTraceInputReader {
    public:
        /// 'traceDir' specifies the operator trace directory.
        OperatorTraceInputReader(
            std::string traceDir,
            RowTypePtr dataType,
            memory::MemoryPool *pool);

        /// Reads from 'dataStream_' and deserializes to 'batch'. Returns false if
        /// reaches to end of the stream and 'batch' is set to nullptr.
        bool read(RowVectorPtr &batch) const;

    private:
        std::unique_ptr<common::FileInputStream> getInputStream() const;

        const std::string traceDir_;
        const serializer::presto::PrestoVectorSerde::PrestoOptions readOptions_{
            true,
            common::CompressionKind_ZSTD, // TODO: Use trace config.
            0.8,
            /*_nullsFirst=*/true
        };
        const std::shared_ptr<filesystems::FileSystem> fs_;
        const RowTypePtr dataType_;
        memory::MemoryPool *const pool_;
        VectorSerde *const serde_;
        const std::unique_ptr<common::FileInputStream> inputStream_;
    };

    /// Used to read an operator trace summary.
    class OperatorTraceSummaryReader {
    public:
        /// 'traceDir' specifies the operator trace directory.
        OperatorTraceSummaryReader(std::string traceDir, memory::MemoryPool *pool);

        /// Read and return the operator trace 'summary'. The function throws if it
        /// fails.
        OperatorTraceSummary read() const;

    private:
        const std::string traceDir_;
        const std::shared_ptr<filesystems::FileSystem> fs_;
        const std::unique_ptr<ReadFile> summaryFile_;
    };

    /// Used to load the input splits from a set of traced 'TableScan' operators for
    /// replay.
    ///
    /// Currently, it only works with 'HiveConnectorSplit'. In the future, it will
    /// be extended to handle more types of splits, such as
    /// 'IcebergHiveConnectorSplit'.
    class OperatorTraceSplitReader {
    public:
        /// 'traceDirs' provides a list of directories with each one containing the
        /// traced split info file for one table scan operator.
        explicit OperatorTraceSplitReader(
            std::vector<std::string> traceDirs,
            memory::MemoryPool *pool);

        /// Reads and deserializes all the traced split strings.
        std::vector<std::string> read() const;

    private:
        static std::vector<std::string> deserialize(common::FileInputStream *stream);

        std::unique_ptr<common::FileInputStream> getSplitInputStream(
            const std::string &traceDir) const;

        const std::vector<std::string> traceDirs_;
        const std::shared_ptr<filesystems::FileSystem> fs_;
        memory::MemoryPool *const pool_;
    };
} // namespace kumo::pollux::exec::trace
